// This file is part of Notepad4.
// See License.txt for details about distribution and modification.
//! Lexer for Swift.

#include <cassert>
#include <cstring>

#include <string>
#include <string_view>
#include <vector>

#include "ILexer.h"
#include "Scintilla.h"
#include "SciLexer.h"

#include "WordList.h"
#include "LexAccessor.h"
#include "Accessor.h"
#include "StyleContext.h"
#include "CharacterSet.h"
#include "StringUtils.h"
#include "LexerModule.h"

using namespace Lexilla;

namespace {

struct EscapeSequence {
	int outerState = SCE_SWIFT_DEFAULT;
	int digitsLeft = 0;
	bool brace = false;

	// highlight any character as escape sequence.
	void resetEscapeState(int state) noexcept {
		outerState = state;
		digitsLeft = 1;
		brace = false;
	}
	bool atEscapeEnd(int ch) noexcept {
		--digitsLeft;
		return digitsLeft <= 0 || !IsHexDigit(ch);
	}
};

struct InterpolatedStringState {
	int state;
	int parenCount;
	int delimiterCount;
};

constexpr bool IsSingleLineString(int state) noexcept {
	return state < SCE_SWIFT_TRIPLE_STRING;
}

constexpr bool IsExtendedDelimitedString(int state) noexcept {
	if constexpr (SCE_SWIFT_STRING_ED & 1) {
		return state & true;
	} else {
		return (state & 1) == 0;
	}
}

enum {
	SwiftLineStateMaskLineComment = 1,		// line comment
	SwiftLineStateMaskImport = (1 << 1),	// import
	SwiftLineStateMaskInterpolation = 1 << 2,
};

//KeywordIndex++Autogenerated -- start of section automatically generated
enum {
	KeywordIndex_Keyword = 0,
	KeywordIndex_Directive = 1,
	KeywordIndex_Attribute = 2,
	KeywordIndex_Class = 3,
	KeywordIndex_Struct = 4,
	KeywordIndex_Protocol = 5,
	KeywordIndex_Enumeration = 6,
	MaxKeywordSize = 36,
};
//KeywordIndex--Autogenerated -- end of section automatically generated

enum class KeywordType {
	None = SCE_SWIFT_DEFAULT,
	Class = SCE_SWIFT_CLASS,
	Struct = SCE_SWIFT_STRUCT,
	Protocol = SCE_SWIFT_PROTOCOL,
	Enum = SCE_SWIFT_ENUM,
	Function = SCE_SWIFT_FUNCTION_DEFINITION,
	Label = SCE_SWIFT_LABEL,
};

constexpr bool IsSwiftIdentifierChar(int state, int ch) noexcept {
	return (state == SCE_SWIFT_IDENTIFIER_BT) ? (ch >= ' ' && ch != '`') : IsIdentifierCharEx(ch);
}

constexpr bool FollowExpression(int chPrevNonWhite, int stylePrevNonWhite) noexcept {
	return chPrevNonWhite == ')' || chPrevNonWhite == ']'
		|| (stylePrevNonWhite >= SCE_SWIFT_OPERATOR_PF && stylePrevNonWhite < SCE_SWIFT_DIRECTIVE)
		|| IsIdentifierCharEx(chPrevNonWhite) || chPrevNonWhite == '`';
}

constexpr bool IsRegexStart(int chPrevNonWhite, int stylePrevNonWhite) noexcept {
	return stylePrevNonWhite == SCE_SWIFT_WORD || !FollowExpression(chPrevNonWhite, stylePrevNonWhite);
}

constexpr bool IsSpaceEquiv(int state) noexcept {
	return state <= SCE_SWIFT_TASKMARKER;
}

void ColouriseSwiftDoc(Sci_PositionU startPos, Sci_Position lengthDoc, int initStyle, LexerWordList keywordLists, Accessor &styler) {
	int lineStateLineType = 0;
	int commentLevel = 0;	// nested block comment level

	KeywordType kwType = KeywordType::None;
	int chBeforeIdentifier = 0;

	int delimiterCount = 0;	// count of '#'
	std::vector<InterpolatedStringState> nestedState;	// string interpolation "\()"

	int visibleChars = 0;
	int visibleCharsBefore = 0;
	int chBefore = 0;
	int chPrevNonWhite = 0;
	int stylePrevNonWhite = SCE_SWIFT_DEFAULT;
	bool insideRegexRange = false; // inside regex character range []
	EscapeSequence escSeq;

	if (startPos != 0) {
		// backtrack to the line starts expression inside interpolated string literal.
		BacktrackToStart(styler, SwiftLineStateMaskInterpolation, startPos, lengthDoc, initStyle);
	}

	StyleContext sc(startPos, lengthDoc, initStyle, styler);
	if (sc.currentLine > 0) {
		const int lineState = styler.GetLineState(sc.currentLine - 1);
		/*
		2: lineStateLineType
		1: SwiftLineStateMaskInterpolation
		1: unused
		8: commentLevel
		: stringInterpolatorCount
		*/
		commentLevel = (lineState >> 4) & 0xff;
		delimiterCount = lineState >> 12;
	}
	if (startPos != 0 && IsSpaceEquiv(initStyle)) {
		// look back for better regex colouring
		LookbackNonWhite(styler, startPos, SCE_SWIFT_TASKMARKER, chPrevNonWhite, stylePrevNonWhite);
	}

	while (sc.More()) {
		switch (sc.state) {
		case SCE_SWIFT_OPERATOR:
		case SCE_SWIFT_OPERATOR2:
		case SCE_SWIFT_OPERATOR_PF:
			sc.SetState(SCE_SWIFT_DEFAULT);
			break;

		case SCE_SWIFT_NUMBER:
			if (!IsDecimalNumberEx(sc.chPrev, sc.ch, sc.chNext)) {
				sc.SetState(SCE_SWIFT_DEFAULT);
			}
			break;

		case SCE_SWIFT_IDENTIFIER:
		case SCE_SWIFT_IDENTIFIER_BT:
		case SCE_SWIFT_ATTRIBUTE:
		case SCE_SWIFT_DIRECTIVE:
		case SCE_SWIFT_VARIABLE:
			if (!IsSwiftIdentifierChar(sc.state, sc.ch)) {
				if (sc.state == SCE_SWIFT_IDENTIFIER_BT) {
					if (sc.ch == '`') {
						sc.Forward();
					}
				} else if (sc.state == SCE_SWIFT_IDENTIFIER || sc.state == SCE_SWIFT_DIRECTIVE) {
					char s[MaxKeywordSize];
					sc.GetCurrent(s, sizeof(s));
					if (sc.state == SCE_SWIFT_DIRECTIVE) {
						if (!keywordLists[KeywordIndex_Directive].InListPrefixed(s + 1, '(')) {
							sc.ChangeState(SCE_SWIFT_MACRO);
						}
					} else if (keywordLists[KeywordIndex_Keyword].InList(s)) {
						sc.ChangeState(SCE_SWIFT_WORD);
						if (StrEqual(s, "import")) {
							if (visibleChars == sc.LengthCurrent()) {
								lineStateLineType = SwiftLineStateMaskImport;
							}
						} else if (StrEqualsAny(s, "class", "extension", "typealias", "as", "is")) {
							kwType = KeywordType::Class;
						} else if (StrEqual(s, "struct")) {
							kwType = KeywordType::Struct;
						} else if (StrEqual(s, "protocol")) {
							kwType = KeywordType::Protocol;
						} else if (StrEqual(s, "enum")) {
							kwType = KeywordType::Enum;
						} else if (StrEqual(s, "func")) {
							kwType = KeywordType::Function;
						} else if (StrEqualsAny(s, "break", "continue")) {
							kwType = KeywordType::Label;
						}
						if (kwType != KeywordType::None) {
							const int chNext = sc.GetLineNextChar();
							if (!IsIdentifierStartEx(chNext)) {
								kwType = KeywordType::None;
							}
						}
					} else if (keywordLists[KeywordIndex_Class].InList(s)) {
						sc.ChangeState(SCE_SWIFT_CLASS);
					} else if (keywordLists[KeywordIndex_Struct].InList(s)) {
						sc.ChangeState(SCE_SWIFT_STRUCT);
					} else if (keywordLists[KeywordIndex_Protocol].InList(s)) {
						sc.ChangeState(SCE_SWIFT_PROTOCOL);
					} else if (keywordLists[KeywordIndex_Enumeration].InList(s)) {
						sc.ChangeState(SCE_SWIFT_ENUM);
					}
				}
				if (sc.state == SCE_SWIFT_IDENTIFIER || sc.state == SCE_SWIFT_IDENTIFIER_BT) {
					if (sc.ch == ':') {
						if (IsJumpLabelPrevASI(chBefore)) {
							sc.ChangeState(SCE_SWIFT_LABEL);
						}
					} else if (sc.ch != '.') {
						if (kwType != KeywordType::None) {
							sc.ChangeState(static_cast<int>(kwType));
						} else {
							const int chNext = sc.GetDocNextChar(sc.ch == '?');
							if (chNext == '(') {
								sc.ChangeState(SCE_SWIFT_FUNCTION);
							} else if ((chBeforeIdentifier == '<' && (chNext == '>' || chNext == '<'))
								|| (chBeforeIdentifier == '[' && (sc.ch == ']' && AnyOf(sc.chNext, '(', ']')))) {
								// type<type>
								// type<type?>
								// type<type<type>>
								// type<type, type>
								// [type]()
								// [[type]]()
								// class type: protocol, protocol {}
								sc.ChangeState(SCE_SWIFT_CLASS);
							}
						}
					}
				}
				if (sc.state != SCE_SWIFT_WORD && sc.ch != '.') {
					kwType = KeywordType::None;
				}
				stylePrevNonWhite = sc.state;
				sc.SetState(SCE_SWIFT_DEFAULT);
			}
			break;

		case SCE_SWIFT_COMMENTLINE:
		case SCE_SWIFT_COMMENTLINEDOC:
			if (sc.atLineStart) {
				sc.SetState(SCE_SWIFT_DEFAULT);
			} else {
				HighlightTaskMarker(sc, visibleChars, visibleCharsBefore, SCE_SWIFT_TASKMARKER);
			}
			break;

		case SCE_SWIFT_COMMENTBLOCK:
		case SCE_SWIFT_COMMENTBLOCKDOC:
			if (sc.Match('*', '/')) {
				sc.Forward();
				--commentLevel;
				if (commentLevel == 0) {
					sc.ForwardSetState(SCE_SWIFT_DEFAULT);
				}
			} else if (sc.Match('/', '*')) {
				sc.Forward();
				++commentLevel;
			} else if (HighlightTaskMarker(sc, visibleChars, visibleCharsBefore, SCE_SWIFT_TASKMARKER)) {
				continue;
			}
			break;

		case SCE_SWIFT_STRING:
		case SCE_SWIFT_STRING_ED:
		case SCE_SWIFT_TRIPLE_STRING:
		case SCE_SWIFT_TRIPLE_STRING_ED:
			if (sc.atLineStart && IsSingleLineString(sc.state)) {
				sc.SetState(SCE_SWIFT_DEFAULT);
			} else if (sc.ch == '\\') {
				int chNext = sc.chNext;
				if (IsExtendedDelimitedString(sc.state)) {
					if (chNext != '#') {
						break;
					}
					const auto delimiter = GetMatchedDelimiterCountEx(styler, sc.currentPos + 1, '#');
					chNext = delimiter.chNext;
					if (delimiter.count != delimiterCount) {
						break;
					}
				}
				if (!IsEOLChar(chNext)) {
					escSeq.resetEscapeState(sc.state);
					sc.SetState(SCE_SWIFT_ESCAPECHAR);
					sc.Advance(delimiterCount);
					sc.Forward();
					if (chNext == '(') {
						sc.ChangeState(SCE_SWIFT_OPERATOR2);
						nestedState.push_back({escSeq.outerState, 1, delimiterCount});
					} else if (sc.Match('u', '{')) {
						escSeq.brace = true;
						escSeq.digitsLeft = 9;
						sc.Forward();
					}
				}
			} else if (sc.ch == '"' && (IsSingleLineString(sc.state) || sc.MatchNext('"', '"'))) {
				if (!IsSingleLineString(sc.state)) {
					sc.Advance(2);
				}
				if (IsExtendedDelimitedString(sc.state)) {
					if (sc.chNext != '#') {
						break;
					}
					const int count = GetMatchedDelimiterCount(styler, sc.currentPos + 1, '#');
					if (count != delimiterCount) {
						break;
					}
					delimiterCount = 0;
					sc.Advance(count);
				}
				sc.Forward();
				if (sc.state == SCE_SWIFT_STRING && (chBefore == '[' || chBefore == ',')) {
					const int chNext = sc.GetLineNextChar();
					if (chNext == ':') {
						sc.ChangeState(SCE_SWIFT_KEY);
					}
				}
				sc.SetState(SCE_SWIFT_DEFAULT);
			}
			break;

		case SCE_SWIFT_REGEX:
		case SCE_SWIFT_REGEX_ED:
			if (sc.atLineStart && sc.state == SCE_SWIFT_REGEX) {
				sc.SetState(SCE_SWIFT_DEFAULT);
			} else if (sc.ch == '\\') {
				sc.Forward();
			} else if (sc.ch == '[' || sc.ch == ']') {
				insideRegexRange = sc.ch == '[';
			} else if (sc.ch == '/' && !insideRegexRange && (sc.state == SCE_SWIFT_REGEX || sc.chNext == '#')) {
				if (sc.state == SCE_SWIFT_REGEX_ED) {
					const int count = GetMatchedDelimiterCount(styler, sc.currentPos + 1, '#');
					if (count != delimiterCount) {
						break;
					}
					delimiterCount = 0;
					sc.Advance(count);
				}
				sc.ForwardSetState(SCE_SWIFT_DEFAULT);
			}
			break;

		case SCE_SWIFT_ESCAPECHAR:
			if (escSeq.atEscapeEnd(sc.ch)) {
				if (escSeq.brace && sc.ch == '}') {
					sc.Forward();
				}
				sc.SetState(escSeq.outerState);
				continue;
			}
			break;
		}

		if (sc.state == SCE_SWIFT_DEFAULT) {
			if (sc.ch == '/') {
				if (sc.chNext == '/' || sc.chNext == '*') {
					visibleCharsBefore = visibleChars;
					const int chNext = sc.chNext;
					sc.SetState((chNext == '/') ? SCE_SWIFT_COMMENTLINE : SCE_SWIFT_COMMENTBLOCK);
					sc.Forward(2);
					if (sc.ch == ':' || sc.ch == '!' || (sc.ch == chNext && sc.chNext != chNext)) {
						static_assert(SCE_SWIFT_COMMENTLINEDOC - SCE_SWIFT_COMMENTLINE == SCE_SWIFT_COMMENTBLOCKDOC - SCE_SWIFT_COMMENTBLOCK);
						sc.ChangeState(sc.state + SCE_SWIFT_COMMENTLINEDOC - SCE_SWIFT_COMMENTLINE);
					}
					if (chNext == '/') {
						if (visibleChars == 0) {
							lineStateLineType = SwiftLineStateMaskLineComment;
						}
					} else {
						commentLevel = 1;
					}
					continue;
				}
				if (sc.chNext > ' ' && IsRegexStart(chPrevNonWhite, stylePrevNonWhite)) {
					insideRegexRange = false;
					sc.SetState(SCE_SWIFT_REGEX);
				} else {
					sc.SetState(SCE_SWIFT_OPERATOR);
				}
			}
			else if (sc.ch == '"') {
				delimiterCount = 0;
				chBefore = chPrevNonWhite;
				sc.SetState(SCE_SWIFT_STRING);
				if (sc.MatchNext('"', '"')) {
					sc.ChangeState(SCE_SWIFT_TRIPLE_STRING);
					sc.Advance(2);
				}
			} else if (IsADigit(sc.ch)) {
				sc.SetState(SCE_SWIFT_NUMBER);
			} else if ((sc.ch == '@' || sc.ch == '#') && IsIdentifierStartEx(sc.chNext)) {
				sc.SetState((sc.ch == '@') ? SCE_SWIFT_ATTRIBUTE : SCE_SWIFT_DIRECTIVE);
			} else if (sc.ch == '$' && IsIdentifierCharEx(sc.chNext)) {
				sc.SetState(SCE_SWIFT_VARIABLE); // closure parameter, synthesized property
			} else if (sc.ch == '#') {
				const auto [count, chNext] = GetMatchedDelimiterCountEx(styler, sc.currentPos, '#');
				if (chNext == '\"' || chNext == '/') {
					insideRegexRange = false;
					delimiterCount = count;
					sc.SetState((chNext == '/') ? SCE_SWIFT_REGEX_ED : SCE_SWIFT_STRING_ED);
					sc.Advance(count);
					if (chNext != '/' && sc.Match('"', '"', '"')) {
						sc.ChangeState(SCE_SWIFT_TRIPLE_STRING_ED);
						sc.Advance(2);
					}
				}
			} else if (IsIdentifierStartEx(sc.ch) || (sc.ch == '`' && sc.chNext >= ' ')) {
				chBefore = chPrevNonWhite;
				if (chPrevNonWhite != '.') {
					chBeforeIdentifier = chPrevNonWhite;
				}
				sc.SetState((sc.ch == '`') ? SCE_SWIFT_IDENTIFIER_BT : SCE_SWIFT_IDENTIFIER);
			} else if (sc.ch == '+' || sc.ch == '-') {
				sc.SetState(SCE_SWIFT_OPERATOR);
				if (sc.ch == sc.chNext) {
					// highlight ++ and -- as different style to simplify regex detection.
					sc.ChangeState(SCE_SWIFT_OPERATOR_PF);
					sc.Forward();
				}
			} else if (IsAGraphic(sc.ch)) {
				const bool interpolating = !nestedState.empty();
				sc.SetState(interpolating ? SCE_SWIFT_OPERATOR2 : SCE_SWIFT_OPERATOR);
				if (interpolating && AnyOf<'(', ')'>(sc.ch)) {
					InterpolatedStringState &state = nestedState.back();
					state.parenCount += '(' + ')' - 2*sc.ch;
					if (state.parenCount <= 0) {
						escSeq.outerState = state.state;
						delimiterCount = state.delimiterCount;
						nestedState.pop_back();
						sc.ForwardSetState(escSeq.outerState);
						continue;
					}
				}
			}
		}

		if (!isspacechar(sc.ch)) {
			visibleChars++;
			if (!IsSpaceEquiv(sc.state)) {
				chPrevNonWhite = sc.ch;
				stylePrevNonWhite = sc.state;
			}
		}
		if (sc.atLineEnd) {
			int lineState = (commentLevel << 4) | (delimiterCount << 12) | lineStateLineType;
			if (!nestedState.empty()) {
				lineState |= SwiftLineStateMaskInterpolation;
			}
			styler.SetLineState(sc.currentLine, lineState);
			lineStateLineType = 0;
			visibleChars = 0;
			visibleCharsBefore = 0;
			kwType = KeywordType::None;
		}
		sc.Forward();
	}

	sc.Complete();
}

struct FoldLineState {
	int lineComment;
	int packageImport;
	constexpr explicit FoldLineState(int lineState) noexcept:
		lineComment(lineState & SwiftLineStateMaskLineComment),
		packageImport((lineState >> 1) & 1) {
	}
};

void FoldSwiftDoc(Sci_PositionU startPos, Sci_Position lengthDoc, int initStyle, LexerWordList /*keywordLists*/, Accessor &styler) {
	const Sci_PositionU endPos = startPos + lengthDoc;
	Sci_Line lineCurrent = styler.GetLine(startPos);
	FoldLineState foldPrev(0);
	int levelCurrent = SC_FOLDLEVELBASE;
	if (lineCurrent > 0) {
		levelCurrent = styler.LevelAt(lineCurrent - 1) >> 16;
		foldPrev = FoldLineState(styler.GetLineState(lineCurrent - 1));
		const Sci_PositionU bracePos = CheckBraceOnNextLine(styler, lineCurrent - 1, SCE_SWIFT_OPERATOR, SCE_SWIFT_TASKMARKER, SCE_SWIFT_DIRECTIVE);
		if (bracePos) {
			startPos = bracePos + 1; // skip the brace
		}
	}

	int levelNext = levelCurrent;
	FoldLineState foldCurrent(styler.GetLineState(lineCurrent));
	Sci_PositionU lineStartNext = styler.LineStart(lineCurrent + 1);

	char chNext = styler[startPos];
	int styleNext = styler.StyleIndexAt(startPos);
	int style = initStyle;
	int visibleChars = 0;

	while (startPos < endPos) {
		const char ch = chNext;
		const int stylePrev = style;
		style = styleNext;
		chNext = styler[++startPos];
		styleNext = styler.StyleIndexAt(startPos);

		switch (style) {
		case SCE_SWIFT_COMMENTBLOCK:
		case SCE_SWIFT_COMMENTBLOCKDOC: {
			const int level = (ch == '/' && chNext == '*') ? 1 : ((ch == '*' && chNext == '/') ? -1 : 0);
			if (level != 0) {
				levelNext += level;
				startPos++;
				chNext = styler[startPos];
				styleNext = styler.StyleIndexAt(startPos);
			}
		} break;

		case SCE_SWIFT_TRIPLE_STRING:
		case SCE_SWIFT_TRIPLE_STRING_ED:
		case SCE_SWIFT_REGEX_ED:
			if (style != stylePrev) {
				levelNext++;
			}
			if (style != styleNext) {
				levelNext--;
			}
			break;

		case SCE_SWIFT_OPERATOR:
		case SCE_SWIFT_OPERATOR2:
			if (ch == '{' || ch == '[' || ch == '(') {
				levelNext++;
			} else if (ch == '}' || ch == ']' || ch == ')') {
				levelNext--;
			}
			break;

		case SCE_SWIFT_DIRECTIVE:
			if (ch == '#') {
				if (chNext == 'i' && styler[startPos + 1] == 'f') {
					levelNext++;
				} else if (chNext == 'e' && styler.Match(startPos, "endif")) {
					levelNext--;
				}
			}
			break;
		}

		if (visibleChars == 0 && !IsSpaceEquiv(style)) {
			++visibleChars;
		}
		if (startPos == lineStartNext) {
			const FoldLineState foldNext(styler.GetLineState(lineCurrent + 1));
			levelNext = sci::max(levelNext, SC_FOLDLEVELBASE);
			if (foldCurrent.lineComment) {
				levelNext += foldNext.lineComment - foldPrev.lineComment;
			} else if (foldCurrent.packageImport) {
				levelNext += foldNext.packageImport - foldPrev.packageImport;
			} else if (visibleChars) {
				const Sci_PositionU bracePos = CheckBraceOnNextLine(styler, lineCurrent, SCE_SWIFT_OPERATOR, SCE_SWIFT_TASKMARKER, SCE_SWIFT_DIRECTIVE);
				if (bracePos) {
					levelNext++;
					startPos = bracePos + 1; // skip the brace
					style = SCE_SWIFT_OPERATOR;
					chNext = styler[startPos];
					styleNext = styler.StyleIndexAt(startPos);
				}
			}

			const int levelUse = levelCurrent;
			int lev = levelUse | (levelNext << 16);
			if (levelUse < levelNext) {
				lev |= SC_FOLDLEVELHEADERFLAG;
			}
			styler.SetLevel(lineCurrent, lev);

			lineCurrent++;
			lineStartNext = styler.LineStart(lineCurrent + 1);
			levelCurrent = levelNext;
			foldPrev = foldCurrent;
			foldCurrent = foldNext;
			visibleChars = 0;
		}
	}
}

}

extern const LexerModule lmSwift(SCLEX_SWIFT, ColouriseSwiftDoc, "swift", FoldSwiftDoc);
